clear all
*cap log close 
set more off
set seed 1 


/*******************************************************************************

BOOTSTRAP VERSION

- Marginal Treatment Effects figure that estimates the MTE *across* beat-shift
fixed effects. While the beat shift fixed effects are controlled for, since the
level of officer randomization is within FEs, the propensity score is not
interacted with beat-shift FE.

- Will also calculate ATT/ATU/ATE by integrating over the MTE

*******************************************************************************/


*** create a dataset that saves all the bootstrap values
set obs 1
gen empty = 1
save "${est}/mte_bootstrap", replace
clear

** set randomization for bootstrap
*set seed 1
*cd "/Users/Felipe/Dropbox/deterrence/tests_new/att"

* read in data ----------
use "${data}/out/4-main", clear

************************************
************************************
// Setup data ------
gen Y = cite_ny1 
gen D = harsh 
gen W = covbin1
egen aggfe = group(troop)
keep citationid Y D Z W officerid totfe

** for the linkage with bootstrap samples
gen obsID = _n

** count the number of observations
count
global N = `r(N)'

compress
tempfile base
save    `base'
clear
************************************
************************************

global NI = 101

///////////////////////////////////////////////////////
////// STEP 2: CREATE RANDOM SAMPLES AND ITERATE /////////

forval i = 1/$NI{

// local i = 1
di "Iteration #`i'"

timer clear 1
timer on 1

use `base', clear
// set obs $N

// local i = 1
** create set of numbers for bootstrap randomization
** first run-through will be the original data
// cap drop baseline weight
if `i' == 1{
// 	gen obsID    = _n
	gen baseline = 1
	gen weight   = 1
}
else{
// 	gen obsID    = floor(runiform()*$N + 1)
	gen baseline = 0
	egen offtag = tag(officerid)
	gen temp    = runiform() if offtag == 1
	bysort officerid: egen weight = max(temp)
	drop offtag temp
}

// ** merge the random set of numbers to the data
// merge m:1 obsID using `base', keep(3)

** calculate propensity score
reghdfe D Z [aw = weight], absorb(totfe, savefe) res
predict pscore1, xbd

** require pscore btw [0, 1]
replace pscore1 = 1 if pscore1>1 & ~mi(pscore1)
replace pscore1 = 0 if pscore1<0

forval k = 2/8{
	gen pscore`k'    = pscore1^`k'
}

** construction of atu/att: integrate over values of u
sort citationid
gen n = _n     if _n<=99
gen u = _n/100 if _n<=99

** ATT/ATU weights over u_D (doesn't depend on MTE estimation)
gen atu_w = .
gen att_w = .
forval j = 1/99{
	** ut_u is indicator for each individual if their p score is above a certain u
	gen ut_u      = pscore1<u[`j']
	qui sum ut_u
	replace atu_w = `r(mean)'     if _n == `j'
	replace att_w = 1 - `r(mean)' if _n == `j'
	drop ut_u
}

// ** loop through possible mtes: go from 1 to 6-order polynomial

forval j = 1/6{

global ppoly ""
local m = `j'+1
forval k = 1/`m' {
	global ppoly "$ppoly pscore`k'"
}

reghdfe Y $ppoly [aw = weight] , absorb(totfe) vce(cluster officerid)

cap drop mte_u_`j' atu_`j' att_`j'
gen mte_u_`j' = 0
gen atu_`j'   = 0
gen att_`j'   = 0
gen ate_`j'   = 0
// local m_ = `m' - 1
forval k = 0/`j'{
	local kp = `k' + 1
	replace mte_u_`j' = mte_u_`j' + `kp'*_b[pscore`kp']*u^`k'
}

sum mte_u_`j' [w = atu_w]
replace atu_`j' = `r(mean)'

sum mte_u_`j' [w = att_w]
replace att_`j' = `r(mean)'

sum mte_u_`j'
replace ate_`j' = `r(mean)'

}

************************
************************
** non parametric reg

preserve
cap rm temp.dta
sort pscore1
binsreg Y pscore1 [aw = weight], absorb(totfe) nbins(12) savedata(temp)

use temp, clear
gen dx = dots_x - dots_x[_n-1]
gen dy = dots_fit - dots_fit[_n-1]
gen dydx = dy/dx 

gen mid = dots_x[_n-1] + (dots_x - dots_x[_n-1])/2
gen x       = mid
gen mtenp_x = dydx 

sort dots_x
set obs 99
gen mte_u_np = .
gen u        = _n/100
gen n        = _n

** map non param mte into u = [0.01 , 0.99]
forval k = 1/99{
	di `k'
	sort n
	sum dots_x
	if `k'/99<`r(min)'{
		replace mte_u_np = mtenp_x[2] if n == `k'
	}
	else if `k'/99>`r(max)'{
		replace mte_u_np = mtenp_x[10] if n == `k'
	}
	else{
		sum n if dots_x>`k'/99
		replace mte_u_np = mtenp_x[`r(min)'] if n == `k'
	}
}

keep mte_u_np u n

tempfile nonparam
save 	`nonparam'

restore
************************
************************

** merge non param mte back to polynomial mtes
merge m:1 n using `nonparam', nogen

** create nonparam treatment effect estimates
sum mte_u_np [w = atu_w]
gen atu_np = `r(mean)'

sum mte_u_np [w = att_w]
gen att_np = `r(mean)'

sum mte_u_np
gen ate_np = `r(mean)'

************************
************************


************************
************************
** append to dataset
preserve
** keep only variables we need
* declare which iteration
// local i = 1
gen it = `i'
label variable it "iteration number"

keep if ~mi(u)
keep it u mte_u_1 - ate_np

append using "${est}/mte_bootstrap" 
save         "${est}/mte_bootstrap" , replace
restore
************************

timer off  1
timer list 1

di "iteration `i' of $NI completed"

}

rm temp.dta 


